#!/bin/sh

. "$SPEC_PATH/abstract/sharded-dbops"

e2e_test_extra_hash() {
  printf '%s\n' E2E_CITUS_POSTGRES_VERSION="$E2E_CITUS_POSTGRES_VERSION"
  "$SHELL" "$PROJECT_PATH/stackgres-k8s/ci/build/build-functions.sh" path_hash \
    "$(realpath --relative-to "$PROJECT_PATH" "$SPEC_PATH/abstract/sharded-dbops")"
}

e2e_test_install() {
  E2E_POSTGRES_VERSION="${E2E_CITUS_POSTGRES_VERSION:-$E2E_POSTGRES_VERSION}"
  CLUSTER_NAME="$(get_sgshardedcluster_name "$SPEC_NAME")"

  create_or_replace_sharded_cluster "$CLUSTER_NAME" "$CLUSTER_NAMESPACE" "3" "2" \
    --set-string "configurations.postgresconfig.postgresql\.conf.citus\.desired_percent_disk_available_after_move=0"

  wait_pods_running "$CLUSTER_NAMESPACE" 6
  wait_sharded_cluster "$CLUSTER_NAME" "$CLUSTER_NAMESPACE"
  switch_sharded_cluster_to_first "$CLUSTER_NAME" "$CLUSTER_NAMESPACE"
}

e2e_test() {
  run_test "Checking that resharding sharded dbops is working" check_sharded_dbops_resharding_is_working
}

check_sharded_dbops_resharding_is_working() {
  SHARDED_DBOPS_NAME="$(get_sgshardeddbops_name resharding-"${CLUSTER_NAME}-$(shuf -i 0-65535 -n 1)")"

  kubectl exec -n "$CLUSTER_NAMESPACE" "${CLUSTER_NAME}-coord-0" -c "postgres-util" -- \
    psql -q -d citus -v ON_ERROR_STOP=1 \
    -c "CREATE TABLE distributed AS SELECT i, (SELECT string_agg(md5(random()::text), ' ') FROM generate_series(1, (30 * random())::integer)) AS t FROM generate_series(1, 100000) AS i" \
    -c "ALTER TABLE distributed ADD PRIMARY KEY (i)"
  wait_until kubectl exec -n "$CLUSTER_NAMESPACE" "${CLUSTER_NAME}-coord-0" -c "postgres-util" -- \
    psql -q -d citus -v ON_ERROR_STOP=1 \
    -c "SELECT create_distributed_table('distributed', 'i')"
  NODE_NAME="$(kubectl exec -n "$CLUSTER_NAMESPACE" "${CLUSTER_NAME}-coord-0" -c "postgres-util" -- \
    psql -q -t -A -d citus -c "SELECT nodename from pg_dist_node WHERE shouldhaveshards LIMIT 1")"
  kubectl exec -n "$CLUSTER_NAMESPACE" "${CLUSTER_NAME}-coord-0" -c "postgres-util" -- \
    psql -q -d citus -v ON_ERROR_STOP=1 \
    -c "SELECT citus_drain_node('$NODE_NAME', ${POSTGRES_REPLICATION_PORT})" \
    -c "SELECT citus_remove_node('$NODE_NAME', ${POSTGRES_REPLICATION_PORT})"
  wait_until kubectl exec -n "$CLUSTER_NAMESPACE" "${CLUSTER_NAME}-coord-0" -c "postgres-util" -- \
      psql -q -d citus -v ON_ERROR_STOP=1 \
      -c "SELECT citus_add_node('$NODE_NAME', ${POSTGRES_REPLICATION_PORT})"
  if ! kubectl exec -n "$CLUSTER_NAMESPACE" "${CLUSTER_NAME}-coord-0" -c "postgres-util" -- \
    psql -q -t -A -d citus -c "SELECT COUNT(*) FROM get_rebalance_table_shards_plan()" | grep -qxF 0
  then
    success "The sharded cluster requires resharding"
  else
    fail "The sharded cluster does not requires resharding"
  fi

  cat << EOF | kubectl create -f -
apiVersion: stackgres.io/v1
kind: SGShardedDbOps
metadata:
  name: $SHARDED_DBOPS_NAME
  namespace: $CLUSTER_NAMESPACE
spec:
  sgShardedCluster: $CLUSTER_NAME
  op: resharding
  resharding:
    citus:
      threshold: 0.1
      drainOnly: false
      rebalanceStrategy: by_disk_size
EOF

  wait_sharded_dbops_is_completed "$SHARDED_DBOPS_NAME"

  TOTAL_SHARDS_SIZE="$(kubectl exec -n "$CLUSTER_NAMESPACE" "${CLUSTER_NAME}-coord-0" -c "postgres-util" -- \
    psql -q -t -A -d citus -c "SELECT SUM(shard_size) FROM citus_shards")"
  if kubectl exec -n "$CLUSTER_NAMESPACE" "${CLUSTER_NAME}-coord-0" -c "postgres-util" -- \
    psql -q -t -A -d citus -c "SELECT SUM(shard_size) / $TOTAL_SHARD_SIZE < 0.3 FROM citus_shards GROUP BY nodename" | grep -qxF t
  then
    fail "The sharded cluster has not been resharded"
  else
    success "The sharded cluster has been resharded"
  fi
}
